import Core.MySQLDB as MySQLDB
import Core.Gadget as Gadget
import pandas as pd


class MAADB(MySQLDB.MySQLDB):

    def __init__(self, address, port, username, password, dbname="maa"):
        super().__init__(address, port, username=username, password=password)
        self.dbname = dbname  # "maa"

    def Get_Simulatefof_Daily_Bar_Dataframe(self, symbol, instrument_type, projection=[], datetime1=None, datetime2=None):
        # 自建模拟组合

        projection = ["date", "net_val", "acc_net_val"]
        filter = [("portfolio_code", symbol), ("subject_type", 0), ("subject_code", '0')]

        if datetime1:
            filter.append(("date", ">=", datetime1))
        if datetime2:
            filter.append(("date", "<=", datetime2))
        sort = [("date", 1)]

        dailybars = self.Find(self.dbname, "maa_portfolio_invest_daily_data", filter=filter, projection=projection,
                              sort=sort)

        if len(dailybars) == 0:
            return pd.DataFrame()

        df_bars = Gadget.DocumentsToDataFrame(dailybars)
        df_bars["date"] = pd.to_datetime(df_bars["date"])
        df_bars["symbol"] = symbol
        return df_bars

    def Get_Selfprod_Daily_Bar_Dataframe(self, symbol, instrument_type, projection=[], datetime1=None, datetime2=None):
        # 判断Symbol是否为数字
        is_valid_symbol = symbol.isdigit()
        if not is_valid_symbol:  # 不是数字，无效代码
            return pd.DataFrame()

        # 自建产品
        projection = ["date", "net_val", "acc_net_val"]
        portfolio_code = self.Find(self.dbname, "maa_products_info", filter=[("id", int(symbol))],\
                         projection=["fof_portfolio_code"])
        #
        if len(portfolio_code) == 0:
            return pd.DataFrame()
        if portfolio_code[0]["fof_portfolio_code"] == None:
            return pd.DataFrame()
        #
        real_portfolio_code = portfolio_code[0]["fof_portfolio_code"].encode('utf-8').decode('utf-8')
        filter = [("portfolio_code", real_portfolio_code), ("subject_type", 0), ("subject_code", '0')]

        if datetime1:
            filter.append(("date", ">=", datetime1))
        if datetime2:
            filter.append(("date", "<=", datetime2))
        sort = [("date", 1)]

        dailybars = self.Find(self.dbname, "maa_portfolio_invest_daily_data", filter=filter, projection=projection,
                              sort=sort)

        if len(dailybars) == 0:
            return pd.DataFrame()

        df_bars = Gadget.DocumentsToDataFrame(dailybars)
        df_bars["date"] = pd.to_datetime(df_bars["date"])
        df_bars["symbol"] = symbol
        return df_bars

